// Copyright (c) OpenMMLab. All rights reserved
#ifndef NMS_CUDA_KERNEL_CUH
#define NMS_CUDA_KERNEL_CUH

#include <float.h>
#include "common_cuda_helper.hpp"
#include "pytorch_cuda_helper.hpp"
#include<iostream>
// 该语句定义了一个名为threadsPerBlock的常量，其值等于unsigned long long int类型的字节数乘以8。
// 这意味着它计算了一个线程块中可以容纳的线程数量。
int const threadsPerBlock = sizeof(unsigned long long int) * 8;

__device__ inline bool devIoU(float const *const a, float const *const b,
                              const int offset, const float threshold) {
  float left = fmaxf(a[0], b[0]), right = fminf(a[2], b[2]);
  float top = fmaxf(a[1], b[1]), bottom = fminf(a[3], b[3]);
  float width = fmaxf(right - left + offset, 0.f),
        height = fmaxf(bottom - top + offset, 0.f);
  float interS = width * height;
  float Sa = (a[2] - a[0] + offset) * (a[3] - a[1] + offset);
  float Sb = (b[2] - b[0] + offset) * (b[3] - b[1] + offset);
  return interS > threshold * (Sa + Sb - interS);
}
// 实现非最大抑制（NMS）算法的CUDA内核函数
// 用于对一系列边界框进行筛选，去除重叠度较高的框
// 参数说明：
// - n_boxes: 输入边界框的数量
// - iou_threshold: 用于判断两个框是否重叠的阈值
// - offset: 边界框的偏移量，用于调整框的位置
// - dev_boxes: 设备（GPU）上的边界框数据
// - dev_mask: 设备（GPU）上的掩码数据，用于记录每个边界框的抑制结果
__global__ static void nms_cuda(const int n_boxes, const float iou_threshold,
                                const int offset, const float *dev_boxes,
                                unsigned long long *dev_mask) {
  // 计算所需的块数量，以适应输入的边界框数量
  int blocks = (n_boxes + threadsPerBlock - 1) / threadsPerBlock;

  // 使用2D循环遍历所有块，处理NMS算法
  CUDA_2D_KERNEL_BLOCK_LOOP(col_start, blocks, row_start, blocks) {
    // 获取当前线程在块内的索引
    const int tid = threadIdx.x;
    /*
        这段代码可能用于对称矩阵操作或需要避免重复计算的场景。若row_start大于col_start，
        意味着处理的是矩阵中的上三角部分（对于下三角矩阵而言）。为了避免重复计算
        （比如只计算下三角或主对角线元素），当检测到row_start大于col_start时，代码提前返回，
        跳过这部分计算。这样可以确保每个相关元素仅被计算一次。
    */
    // 确保行号不大于列号，以避免重复计算 blockidx.y>blockidx.x,return
    if (row_start > col_start) return;

    // 计算当前块的行和列应处理的边界框数量
    // n_boxes-blockIdx.y*threadsPerBlock
    // row_start * threadsPerBlock 计算出当前行之前已经处理的边界框数量。
    // n_boxes - row_start * threadsPerBlock 则得到当前行及后续可能还需要处理的边界框数量。
    //然而，由于每个线程块最多只能处理 threadsPerBlock 个边界框，所以需要使用 fminf 函数取这两个值中的较小值，
    //作为当前行实际能够处理的边界框数量，并将其存储在 row_size 中。
    const int row_size =
        fminf(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
    const int col_size =
        fminf(n_boxes - col_start * threadsPerBlock, threadsPerBlock);

    // 共享内存中存储当前块所需的边界框数据
    __shared__ float block_boxes[threadsPerBlock * 4];
    // 每个线程负责将全局内存中的边界框数据加载到共享内存中
    // (threadsPerBlock * col_start + tid) 计算出当前线程对应的全局边界框数据的索引。
    if (tid < col_size) {
      block_boxes[tid * 4 + 0] =
          dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 0];
      block_boxes[tid * 4 + 1] =
          dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 1];
      block_boxes[tid * 4 + 2] =
          dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 2];
      block_boxes[tid * 4 + 3] =
          dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 3];
    }
    // 确保所有线程都完成了数据加载
    __syncthreads();

    // 每个线程处理不同的行边界框
    if (tid < row_size) {
      // 获取当前线程应处理的边界框索引
      const int cur_box_idx = threadsPerBlock * row_start + tid;
      // 指向当前边界框的数据
      const float *cur_box = dev_boxes + cur_box_idx * 4;
      // 初始化相关变量
      int i = 0;
      unsigned long long int t = 0;
      int start = 0;
      // 如果行号等于列号，从下一个边界框开始处理
      if (row_start == col_start) {
        start = tid + 1;
      }
      // 遍历列边界框，计算重叠度并更新掩码
      for (i = start; i < col_size; i++) {
        // 如果当前边界框与其它边界框的重叠度超过阈值，则在掩码中做标记
        if (devIoU(cur_box, block_boxes + i * 4, offset, iou_threshold)) {
          t |= 1ULL << i;
        }
      }
      // 将计算结果存储到掩码数组中
      dev_mask[cur_box_idx * gridDim.y + col_start] = t;
    }
  }
}

// 使用CUDA __global__关键字声明一个全局函数，用于从掩码中收集并保留特定的框。
// 这个函数旨在处理一组框（如边界框），根据给定的掩码确定哪些框应该被保留。
// 参数:
// - keep: 一个指针，用于标记哪些框应该被保留。
// - dev_mask: 一个指向设备（GPU）上掩码数组的指针，该数组指示哪些框应该被保留。
// - n_boxes: 表示框的数量。
/*
创建变量remv数组（长为num_bbox，用来保存该bbox是否抛弃，若抛弃则为1，即存在iou>阈值），
keep数组（长为num_bbox，最后返回的bool tensor，若保留则为true）。
依次遍历mask矩阵的每一行i：
  若remv[i]为0，则设置keep[i]为true，并且遍历mask的第i行，
    若mask[i][j]为1，则remv[j]也设置为1；
  若remv[i]为1，则keep[i]为false。
*/
__global__ static void gather_keep_from_mask(bool *keep,
                                             const unsigned long long *dev_mask,
                                             const int n_boxes) {
  // 计算列块的数量，用于处理共享内存中的去除标记。
  const int col_blocks = (n_boxes + threadsPerBlock - 1) / threadsPerBlock;
  // 获取线程在块中的索引。
  const int tid = threadIdx.x;

  // 共享内存数组，用于标记哪些框应该被去除。
  // 注意：数组大小在运行时确定，由CUDA运行时分配。
  extern __shared__ unsigned long long removed[];

  // 初始化去除标记数组。
  for (int i = tid; i < col_blocks; i += blockDim.x) {
    removed[i] = 0;
  }
  // 确保所有线程都完成初始化后再继续。
  __syncthreads();

  // 遍历列块，处理每个块中的线程块。
  for (int nblock = 0; nblock < col_blocks; ++nblock) {
    // 获取当前块的去除值。
    auto removed_val = removed[nblock];
    // 确保所有线程都完成当前块的处理后再继续。
    __syncthreads();
    // 计算当前块的起始偏移量。
    const int i_offset = nblock * threadsPerBlock;
    // 使用#pragma unroll进行循环展开，以提高性能。
#pragma unroll
    //每个线程执行64次循环
    for (int inblock = 0; inblock < threadsPerBlock; ++inblock) {
      // 计算当前线程在块中的索引。
      const int i = i_offset + inblock;
      // 确保索引不超过框的数量。
      if (i >= n_boxes) break;
      // 如果当前线程对应的框未被标记为去除，则进行处理。
      if (!(removed_val & (1ULL << inblock))) {
        // 全局只有这一句修改keep的地方，不用担心其他线程修改
        //线程0负责标记该输出框为保留。
        if (tid == 0) {
          keep[i] = true;
        }
        // 获取当前框的掩码地址。
        auto p = dev_mask + i * col_blocks;
        // 去除所有与当前框重叠的框。
        // 遍历mask的第i行，将值为1的bbox抛弃
        for (int j = tid; j < col_blocks; j += blockDim.x) {
          if (j >= nblock) removed[j] |= p[j];
        }
        // 确保所有线程都完成当前块的更新后再继续。
        __syncthreads();
        // 更新当前块的去除值。
        removed_val = removed[nblock];
      }
    }
  }
}
#endif  // NMS_CUDA_KERNEL_CUH